"""
convert ovis opequt to images and json
"""

import datasets
import json


dataset = datasets.load_dataset(
    # "TIGER-Lab/Mantis-Instruct", "multi_vqa", revision="script"
    # "TIGER-Lab/Mantis-Instruct", "llava_665k_multi", revision="script"
    "ovis_data/meta_files/v1_5",
    data_files="pixelprose-14m.parquet",
    # "TIGER-Lab/Mantis-Instruct", data_files="spot-the-diff/train-*-of-00001.parquet"
)
# dataset['train'][0]['images']: processed absolution valid path of the downloaded images on your local machine

results = []

for a in dataset["train"]:
    try:
        print(a)
        images = a["image_url"]
        images_files = []
        for img in images:
            # img_f = 'images/' + img['path'].split('train_images')[-1].lstrip('/')
            # img_f = img['path'].split('train_images')[-1].lstrip('/')
            img_f = img["path"].lstrip("/")
            # img_f = img_f.replace('images/sharegpt4v/llava/', './')
            # img_f = img_f.replace('images/sharegpt4v/', '')
            images_files.append(img_f)

        if len(images_files) == 0:
            print(f"{a} has no images.")
            continue
        data_dict = {
            "id": f"multi_{a['id']}",
            "image": images_files,
        }
        data_dict["conversations"] = []
        for txt in a["conversation"]:
            if txt["role"] == "user":
                data_dict["conversations"].append(
                    {"from": "human", "value": f"{txt['content']}"},
                )
            if txt["role"] == "assistant":
                content = txt["content"]
                if content.startswith("Answer: "):
                    content = content.replace("Answer: ", "")
                data_dict["conversations"].append(
                    {"from": "gpt", "value": content},
                )

        # print(data_dict)

        if len(data_dict["conversations"]) > 5:
            # depart conversations into 2 parts
            for i in range(0, len(data_dict["conversations"]), 10):
                data_dict_i = {}
                data_dict_i["id"] = data_dict["id"] + f"_{i//10}"
                data_dict_i["image"] = data_dict["image"]
                data_dict_i["conversations"] = data_dict["conversations"][i : i + 10]
                if i != 0:
                    data_dict_i["conversations"][0][
                        "value"
                    ] = f"{len(data_dict['image']) * '<image> '}{data_dict_i['conversations'][0]['value']}"

                is_next_contains_tag = False
                for fi, fuck in enumerate(data_dict_i["conversations"]):
                    if fi != 0 and fuck["value"].count("<image>") >= 1:
                        is_next_contains_tag = True
                if not is_next_contains_tag:
                    results.append(data_dict_i)
                else:
                    print(f"passing i: {data_dict_i}")
        else:
            is_next_contains_tag = False
            for m, jj in enumerate(data_dict["conversations"]):
                if m != 0 and jj["value"].count("<image>") >= 1:
                    is_next_contains_tag = True
            if not is_next_contains_tag:
                results.append(data_dict)
            else:
                print(f"passing {data_dict}")
    except Exception as e:
        print(e)
        import traceback

        traceback.print_exc()
        continue

print(f"All samples: {len(results)}")
with open("ovis_pixelprose.json", "w", encoding="utf-8") as f:
    json.dump(results, f, indent=2, ensure_ascii=False)
